import numpy as np
import torch
from torch.utils.data import Dataset
from scipy.io import loadmat


class DeapDataset(Dataset):
    def __init__(self, file_root, start_idx, end_idx):
        self.file_paths = []
        self.label_paths = []
        for i in range(start_idx, end_idx):
            self.file_paths.append(f'{file_root}/s{str(i).zfill(2)}/{i}.mat')
            self.label_paths.append(
                f'C:/Users/john/Desktop/learn-eeg/plv_matrix/s{str(i).zfill(2)}/labels.npy')
        self.samples_per_file = 4600

        self.cached_data = {}
        self.cached_label = {}

    def __len__(self):
        return len(self.file_paths) * self.samples_per_file

    def __getitem__(self, idx):
        file_idx = idx // self.samples_per_file
        sample_idx = idx % self.samples_per_file

        if file_idx not in self.cached_data:
            self.cached_data[file_idx] = loadmat(self.file_paths[file_idx])['preprocessed_sub_4600']
            self.cached_label[file_idx] = np.load(self.label_paths[file_idx])

        label = 1 if self.cached_label[file_idx][sample_idx][0] > 5 else 0

        return torch.tensor(self.cached_data[file_idx][sample_idx], dtype=torch.float32), torch.tensor(label, dtype=torch.long)


if __name__ == '__main__':
    dataset = DeapDataset('C:/Users/john/Desktop/learn-eeg/plv_matrix')
    print(len(dataset))
    print(dataset[0][1])